# include "vd.h"
# include "VD_CreateDevice.h"
# include "VD_DeleteDevice.h"
# include "VD_Crypto.h"

NTSTATUS VD_DeviceControl (
    IN PDEVICE_OBJECT   DeviceObject,
    IN PIRP             Irp
    )
{
    PDEVICE_EXTENSION				device_extension;
    PDEVICE_OBJECT					device_object;
	PIO_STACK_LOCATION				io_stack;
    NTSTATUS						status;
	PCREATE_DISK_INFO				file_info;
	PLONG							Device_Number;
	size_t							TmpLen;			

    device_extension = (PDEVICE_EXTENSION) DeviceObject->DeviceExtension;

    io_stack = IoGetCurrentIrpStackLocation(Irp);

    switch (io_stack->Parameters.DeviceIoControl.IoControlCode)
    {
   
	case IOCTL_VD_CHECK_FILE:
	{
		LONG						PasswordLen;		
		UNICODE_STRING			    tmpW;
		ANSI_STRING					tmpA;	
	
		DbgPrint("IOCTL Check File \n");
		
		if (io_stack->Parameters.DeviceIoControl.InputBufferLength <
							            sizeof(PASSWORD_DATA))
		{
			DbgPrint("Invalid PASSWORD_DATA Buffer \n");
			status = STATUS_INVALID_PARAMETER;
			Irp->IoStatus.Information = 0;
			break;
		}
		

		if (((PPASSWORD_DATA)Irp->AssociatedIrp.SystemBuffer)->PasLen == 0)
		{
			DbgPrint("Can't create file.PasswordLen = 0!!! \n");
			status = STATUS_INVALID_PARAMETER;
			Irp->IoStatus.Information = 0;
			break;
		}

		status =  RtlStringCbLengthW(
									((PPASSWORD_DATA)Irp->AssociatedIrp.SystemBuffer)->Password,
									NTSTRSAFE_MAX_CCH * sizeof(WCHAR),
									&TmpLen
									);
		if(!NT_SUCCESS(status))
		{
			DbgPrint("Can't create file.Can't read Password Name Len!!! \n");				
			status = STATUS_INVALID_PARAMETER;
			Irp->IoStatus.Information = 0;
			break;	
		}

		if (TmpLen == 0)
		{
			DbgPrint("Can't create file.Password Name is empty!!! \n");
			status = STATUS_INVALID_PARAMETER;
			Irp->IoStatus.Information = 0;
			break;
		}

		PasswordLen = ((PPASSWORD_DATA)Irp->AssociatedIrp.SystemBuffer)->PasLen;
		
		tmpW.Length = PasswordLen;
		tmpW.MaximumLength = PasswordLen;
		tmpW.Buffer = (PWCHAR)ExAllocatePoolWithTag(
											NonPagedPool,
											(PasswordLen + sizeof(WCHAR)),
											POOL_TAG
											   );
		if (tmpW.Buffer  == 0)
		{
			DbgPrint("Can't Allocate memory for tmpW!!! \n");
			status = STATUS_INVALID_PARAMETER;
			Irp->IoStatus.Information = 0;
			break;
		}
		
		RtlCopyMemory(
					tmpW.Buffer,
					((PPASSWORD_DATA)Irp->AssociatedIrp.SystemBuffer)->Password,
					PasswordLen
					);

		tmpW.Buffer[PasswordLen/sizeof(WCHAR)] = L'\0';
		
		RtlUnicodeStringToAnsiString(
									&tmpA,
									&tmpW,
									TRUE
									);
		
		device_extension->password = (PUCHAR)ExAllocatePoolWithTag(
																NonPagedPool,
																(tmpA.Length + sizeof(CHAR)),
																POOL_TAG
																  );
		if (device_extension->password  == 0)
		{
			DbgPrint("Can't Allocate memory for password buffer!!! \n");
			status = STATUS_INSUFFICIENT_RESOURCES;
			Irp->IoStatus.Information = 0;
			break;
		}

		RtlCopyMemory(
					device_extension->password,
					tmpA.Buffer,
					tmpA.Length
					);		

		device_extension->password[tmpA.Length] = '\0';
		device_extension->pas_len = tmpA.Length ;

		RtlFreeAnsiString(&tmpA);
		ExFreePoolWithTag(tmpW.Buffer,POOL_TAG);

		IoMarkIrpPending(Irp);
           
		ExInterlockedInsertTailList(
				&device_extension->list_head,
                &Irp->Tail.Overlay.ListEntry,
                &device_extension->list_lock
                );
			
		KeSetEvent(
                &device_extension->request_event,
                (KPRIORITY) 0,
                FALSE
                );
		
		status = STATUS_PENDING;
		break;
    }

	case IOCTL_VD_CREATE_DEVICE:
	{
	if (io_stack->Parameters.DeviceIoControl.InputBufferLength < sizeof(LONG))
    {
		DbgPrint("Invalid CREATE_DEVICE Buffer \n");
		status = STATUS_INVALID_PARAMETER;
        Irp->IoStatus.Information = 0;
        break;
    }	
	
	Device_Number = (LONG*)Irp->AssociatedIrp.SystemBuffer;
	if (DiskDeviceStatus[*Device_Number] == L'0'){
		status = VD_CreateDevice(BaseDriverObject,*Device_Number);
		if (!NT_SUCCESS(status))
			DbgPrint("Unable to create New Device\n");
	}
	else{
		DbgPrint("New device wasn't created. Using existing device \n");
		status = STATUS_SUCCESS;
	}
	break;
	}
	
	case IOCTL_VD_DELETE_DEVICE:
	{
	
	if (io_stack->Parameters.DeviceIoControl.InputBufferLength < sizeof(LONG))
	{
		DbgPrint("Invalid DELETE_DEVICE Buffer \n");
		status = STATUS_INVALID_PARAMETER;
        Irp->IoStatus.Information = 0;
        break;
    }		
	
	Device_Number = (LONG*)Irp->AssociatedIrp.SystemBuffer;
	
	status = VD_DeleteDevice(*Device_Number);
		if (!NT_SUCCESS(status))
			DbgPrint("Unable to delete device \n");
	break;
	}

	case IOCTL_VD_OPEN_FILE:
    {		
	
	WCHAR						device_name_buffer[MAXIMUM_FILENAME_LENGTH];
	UNICODE_STRING			    device_name;
	HANDLE						thread_handle;	
	LONG						FileNameLen;
	LONG						PasswordLen;		
	UNICODE_STRING			    tmpW;
	ANSI_STRING					tmpA;	

	DbgPrint("DeviceControl IOCTL_VD_OPEN_FILE\n");
                      
	if (io_stack->Parameters.DeviceIoControl.InputBufferLength <
							            sizeof(CREATE_DISK_INFO))
	{
		DbgPrint("Can't create file. Invalid CREATE_DISK_INFO Buffer \n");
		status = STATUS_INVALID_PARAMETER;
        Irp->IoStatus.Information = 0;
        break;
	}
	
	if (((PCREATE_DISK_INFO)Irp->AssociatedIrp.SystemBuffer)->FileSize == 0)
	{
		DbgPrint("FileSize Length = 0!!! \n");
		status = STATUS_INVALID_PARAMETER;
        Irp->IoStatus.Information = 0;
		break;
	}
	

	if (((PCREATE_DISK_INFO)Irp->AssociatedIrp.SystemBuffer)->FileNameLength == 0)
	{
		DbgPrint("Can't create file.FileName Length = 0!!! \n");
		status = STATUS_INVALID_PARAMETER;
        Irp->IoStatus.Information = 0;
		break;
	}


	FileNameLen = ((PCREATE_DISK_INFO)Irp->AssociatedIrp.SystemBuffer)->FileNameLength;
	
	device_extension->file_name.Length = FileNameLen;
	device_extension->file_name.MaximumLength = FileNameLen;
	device_extension->file_name.Buffer = (PWCHAR)ExAllocatePoolWithTag(
																	NonPagedPool,
																	(FileNameLen + sizeof(WCHAR)),
																	POOL_TAG
																	);
	if (device_extension->file_name.Buffer == NULL)
    {
		DbgPrint("Unable to allocate memory for CryptedPassword \n");
		Irp->IoStatus.Status = STATUS_INSUFFICIENT_RESOURCES;
		Irp->IoStatus.Information = 0;
		break;    
    }
				
	RtlCopyMemory(
				device_extension->file_name.Buffer,
				((PCREATE_DISK_INFO)Irp->AssociatedIrp.SystemBuffer)->FileName,
				FileNameLen
				);
		
	device_extension->file_size.QuadPart = ((PCREATE_DISK_INFO)Irp->AssociatedIrp.SystemBuffer)->FileSize;
		
	IoMarkIrpPending(Irp);
           
	ExInterlockedInsertTailList(
				&device_extension->list_head,
                &Irp->Tail.Overlay.ListEntry,
                &device_extension->list_lock
                );
			
	KeSetEvent(
                &device_extension->request_event,
                (KPRIORITY) 0,
                FALSE
                );
    status = STATUS_PENDING;
	break;
	}

	case IOCTL_VD_CLOSE_FILE:
    {
	IoMarkIrpPending(Irp);

    ExInterlockedInsertTailList(
                &device_extension->list_head,
                &Irp->Tail.Overlay.ListEntry,
                &device_extension->list_lock
                );
			
    KeSetEvent(
                &device_extension->request_event,
                (KPRIORITY) 0,
                FALSE
                );
    status = STATUS_PENDING;
	break;
    }

    case IOCTL_VD_QUERY_FILE:
    {
	PCREATE_DISK_INFO	create_file_info;
			
	DbgPrint("IOCTL_VD_QUERY_FILE\n");	
            
	if (io_stack->Parameters.DeviceIoControl.OutputBufferLength <
       sizeof(CREATE_DISK_INFO) + device_extension->file_name.Length - sizeof(UCHAR))
    {
		status = STATUS_BUFFER_TOO_SMALL;
        Irp->IoStatus.Information = 0;
        break;
     }

     create_file_info = (PCREATE_DISK_INFO) Irp->AssociatedIrp.SystemBuffer;
     create_file_info->FileSize = device_extension->file_size.QuadPart;
     create_file_info->FileNameLength = device_extension->file_name.Length;

     RtlCopyMemory(
                create_file_info->FileName,
                device_extension->file_name.Buffer,
                device_extension->file_name.Length
                );

     status = STATUS_SUCCESS;
     Irp->IoStatus.Information = sizeof(CREATE_DISK_INFO) +
     create_file_info->FileNameLength - sizeof(UCHAR);
	 break;
     }

    case IOCTL_DISK_CHECK_VERIFY:
    {
    //DbgPrint("IOCTL_DISK_CHECK_VERIFY\n");
	status = STATUS_SUCCESS;
    Irp->IoStatus.Information = 0;
    break;
    }

    case IOCTL_DISK_GET_DRIVE_GEOMETRY:
	case IOCTL_STORAGE_GET_MEDIA_TYPES_EX:
    {
	
	PDISK_GEOMETRY		disk_geometry;
    ULONGLONG			length;
    ULONG				sector_size;
			
    if (io_stack->Parameters.DeviceIoControl.OutputBufferLength <
         sizeof(DISK_GEOMETRY))
    {
		status = STATUS_BUFFER_TOO_SMALL;
		Irp->IoStatus.Information = 0;
		break;
    }

    disk_geometry = (PDISK_GEOMETRY) Irp->AssociatedIrp.SystemBuffer;
	length = device_extension->file_size.QuadPart - PassOffset;
    // DbgPrint("IOCTL_DISK_GET_DRIVE_GEOMETRY FILE_SIZE %d \n", length);
	sector_size = 512;
   
    disk_geometry->Cylinders.QuadPart = length / sector_size / 32 / 2;
    disk_geometry->MediaType = FixedMedia;
    disk_geometry->TracksPerCylinder = 2;
    disk_geometry->SectorsPerTrack = 32;
    disk_geometry->BytesPerSector = sector_size;

    status = STATUS_SUCCESS;
    Irp->IoStatus.Information = sizeof(DISK_GEOMETRY);

    break;
    }

    case IOCTL_DISK_GET_LENGTH_INFO:
    {
           
	PGET_LENGTH_INFORMATION get_length_information;
	// DbgPrint("IOCTL_DISK_GET_LENGTH_INFO\n");
    if (io_stack->Parameters.DeviceIoControl.OutputBufferLength <
									sizeof(GET_LENGTH_INFORMATION))
    {
		status = STATUS_BUFFER_TOO_SMALL;
		Irp->IoStatus.Information = 0;
		break;
    }

    get_length_information = (PGET_LENGTH_INFORMATION) Irp->AssociatedIrp.SystemBuffer;
    get_length_information->Length.QuadPart = device_extension->file_size.QuadPart;
    status = STATUS_SUCCESS;
    Irp->IoStatus.Information = sizeof(GET_LENGTH_INFORMATION);
	break;
    }

    case IOCTL_DISK_GET_PARTITION_INFO:
    {
	
	PPARTITION_INFORMATION  partition_information;
    ULONGLONG               length;
	//DbgPrint("IOCTL_DISK_GET_PARTITION_INFO\n");
    if (io_stack->Parameters.DeviceIoControl.OutputBufferLength <
									sizeof(PARTITION_INFORMATION))
    {
		status = STATUS_BUFFER_TOO_SMALL;
		Irp->IoStatus.Information = 0;
        break;
    }

    partition_information = (PPARTITION_INFORMATION) Irp->AssociatedIrp.SystemBuffer;

	length = device_extension->file_size.QuadPart - - PassOffset;

    partition_information->StartingOffset.QuadPart = PassOffset;
    partition_information->PartitionLength.QuadPart = length;
    partition_information->HiddenSectors = 1;
    partition_information->PartitionNumber = 0;
    partition_information->PartitionType = 0;
    partition_information->BootIndicator = FALSE;
    partition_information->RecognizedPartition = FALSE;
    partition_information->RewritePartition = FALSE;

	status = STATUS_SUCCESS;
    Irp->IoStatus.Information = sizeof(PARTITION_INFORMATION);
	break;
	}

    case IOCTL_DISK_GET_PARTITION_INFO_EX:
    {
				
	PPARTITION_INFORMATION_EX   partition_information_ex;
    ULONGLONG                   length;
	//DbgPrint("IOCTL_DISK_GET_PARTITION_INFO_EX\n");
    if (io_stack->Parameters.DeviceIoControl.OutputBufferLength <
								sizeof(PARTITION_INFORMATION_EX))
    {
		status = STATUS_BUFFER_TOO_SMALL;
        Irp->IoStatus.Information = 0;
        break;
	}

    partition_information_ex = (PPARTITION_INFORMATION_EX) Irp->AssociatedIrp.SystemBuffer;

	length = device_extension->file_size.QuadPart - PassOffset;

    partition_information_ex->PartitionStyle = PARTITION_STYLE_MBR;
    partition_information_ex->StartingOffset.QuadPart = PassOffset;
    partition_information_ex->PartitionLength.QuadPart = length;
    partition_information_ex->PartitionNumber = 0;
    partition_information_ex->RewritePartition = FALSE;
    partition_information_ex->Mbr.PartitionType = 0;
    partition_information_ex->Mbr.BootIndicator = FALSE;
    partition_information_ex->Mbr.RecognizedPartition = FALSE;
    partition_information_ex->Mbr.HiddenSectors = 1;

    status = STATUS_SUCCESS;
    Irp->IoStatus.Information = sizeof(PARTITION_INFORMATION_EX);
    break;
	}

    case IOCTL_DISK_IS_WRITABLE:
	{
	//DbgPrint("IOCTL_DISK_IS_WRITABLE\n");  
	status = STATUS_SUCCESS;
    Irp->IoStatus.Information = 0;
    break;
    }

    case IOCTL_DISK_MEDIA_REMOVAL:
    {
	//DbgPrint("IOCTL_DISK_MEDIA_REMOVAL\n");	
	status = STATUS_SUCCESS;
    Irp->IoStatus.Information = 0;
    break;
	}

    case IOCTL_DISK_SET_PARTITION_INFO:
    {
	//DbgPrint("IOCTL_DISK_SET_PARTITION_INFO\n");
    if (io_stack->Parameters.DeviceIoControl.InputBufferLength <
								sizeof(SET_PARTITION_INFORMATION))
    {
		status = STATUS_INVALID_PARAMETER;
        Irp->IoStatus.Information = 0;
        break;
    }

    status = STATUS_SUCCESS;
    Irp->IoStatus.Information = 0;
	break;
    }

    case IOCTL_DISK_VERIFY:
    {
			
    PVERIFY_INFORMATION verify_information;
	//DbgPrint("IOCTL_DISK_VERIFY\n");
    if (io_stack->Parameters.DeviceIoControl.InputBufferLength <
									sizeof(VERIFY_INFORMATION))
    {
		status = STATUS_INVALID_PARAMETER;
        Irp->IoStatus.Information = 0;
        break;
    }

    verify_information = (PVERIFY_INFORMATION) Irp->AssociatedIrp.SystemBuffer;

    status = STATUS_SUCCESS;
    Irp->IoStatus.Information = verify_information->Length;
	break;
    }

	case IOCTL_STORAGE_GET_HOTPLUG_INFO:
    {
	
	PSTORAGE_HOTPLUG_INFO hotplug_info;

	if (io_stack->Parameters.DeviceIoControl.OutputBufferLength <
									sizeof(STORAGE_HOTPLUG_INFO))
	{
		status = STATUS_INVALID_PARAMETER;
		Irp->IoStatus.Information = 0;
		break;
	}

	hotplug_info = (PSTORAGE_HOTPLUG_INFO)
	Irp->AssociatedIrp.SystemBuffer;

	hotplug_info->Size = sizeof(STORAGE_HOTPLUG_INFO);
	   
	hotplug_info->MediaRemovable = FALSE;
	hotplug_info->MediaHotplug = FALSE;
	hotplug_info->DeviceHotplug = FALSE;
	hotplug_info->WriteCacheEnableOverride = FALSE;
	 
	status = STATUS_SUCCESS;
	Irp->IoStatus.Information = sizeof(STORAGE_HOTPLUG_INFO);
	break;
    }

    default:
    {
    //    DbgPrint("STATUS_INVALID_DEVICE_REQUEST\n");
	status = STATUS_INVALID_DEVICE_REQUEST;
    Irp->IoStatus.Information = 0;
    }
    }

    if (status != STATUS_PENDING)
    {
 
		Irp->IoStatus.Status = status;
        IoCompleteRequest(Irp, IO_NO_INCREMENT);
    }

    return status;
}
